// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/common/base/semaphore.h>
#include <pollux/common/base/simd_util.h>
#include <pollux/common/time/timer.h>

#include <melon/benchmark.h>
#include <melon/executors/cpu_thread_pool_executor.h>
#include <melon/init/init.h>
#include <pollux/flag_definitions/flags.h>

#include <iostream>

TURBO_FLAG(int64_t,repeats, 100, "Number of repetitions");
TURBO_FLAG(int64_t,bytes, 16 * 1024 * 1024, "Bytes to copy in one repetition");
TURBO_FLAG(int64_t,threads, 8, "Number of threads to use within one copy");
TURBO_FLAG(bool,system_memcpy, true, "Use libc memcpy");

using namespace kumo::pollux;

uint64_t sum(uint64_t *data, int32_t size) {
    uint64_t sum = 0;
    for (auto i = 0; i < size; ++i) {
        sum += data[i];
    }
    return sum;
}

struct CopyCallable {
    void *source;
    void *destination;
    int64_t size;
    Semaphore *sem;

    void operator()() {
        if (turbo::get_flag(FLAGS_system_memcpy)) {
            memcpy(destination, source, size);
        } else {
            simd::memcpy(destination, source, size);
        }
        sem->release();
    }
};

int main(int argc, char **argv) {
    constexpr int32_t kAlignment = melon::hardware_destructive_interference_size;
    melon::Init init{&argc, &argv};
    auto chunk = bits::roundUp(
        std::max<int64_t>(turbo::get_flag(FLAGS_bytes) / turbo::get_flag(FLAGS_threads), kAlignment), kAlignment);
    int64_t bytes = chunk * turbo::get_flag(FLAGS_threads);
    auto executor = std::make_unique<melon::CPUThreadPoolExecutor>(turbo::get_flag(FLAGS_threads));
    void *other = aligned_alloc(kAlignment, bytes);
    void *source = aligned_alloc(kAlignment, bytes);
    void *destination = aligned_alloc(kAlignment, bytes);
    // Write all memory once outside of timed section to make them resident.
    memset(other, 1, bytes);
    memset(source, 1, bytes);
    memset(destination, 1, bytes);

    Semaphore sem(0);
    std::vector<CopyCallable> ops;
    ops.resize(turbo::get_flag(FLAGS_threads));
    volatile uint64_t totalSum = 0;
    uint64_t totalUsec = 0;
    for (auto repeat = 0; repeat < turbo::get_flag(FLAGS_repeats); ++repeat) {
        // Read once through 'other' to clear cache effects.
        melon::doNotOptimizeAway(
            totalSum +=
            sum(reinterpret_cast<uint64_t *>(other), bytes / sizeof(int64_t)));
        uint64_t usec = 0; {
            MicrosecondTimer timer(&usec);
            for (auto i = 0; i < turbo::get_flag(FLAGS_threads); ++i) {
                int64_t offset1 = chunk * i;
                ops[i].source = reinterpret_cast<char *>(source) + offset1;
                ops[i].destination = reinterpret_cast<char *>(destination) + offset1;
                ops[i].size = chunk;
                ops[i].sem = &sem;
                executor->add(ops[i]);
            }
            for (auto i = 0; i < turbo::get_flag(FLAGS_threads); ++i) {
                sem.acquire();
            }
        }
        totalUsec += usec;
    }
    std::cout << fmt::format(
                "{} repeats {} bytes {} threads: {} usec {} GB/s",
                turbo::get_flag(FLAGS_repeats),
                bytes,
                turbo::get_flag(FLAGS_threads),
                totalUsec,
                bytes * turbo::get_flag(FLAGS_repeats) / static_cast<float>(1 << 30) /
                (static_cast<float>(totalUsec + 1) / 1000000.0))
            << std::endl;
    free(source);
    free(destination);
    free(other);
    return 0;
}
